# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: yuxuan

import json
import os

from absl import (flags, app)
from coin.strategy.basis_strat.sim.sim_job_generator import \
    run_jobs, get_trading_dates

FLAGS = flags.FLAGS


def _generate_config(param_config, keys, i, config):
  if i < len(keys):
    key = keys[i]
    values = param_config[key]
    for val in values:
      config[key] = val
      if i == len(keys) - 1:
        yield config.copy()
      else:
        for config in _generate_config(param_config, keys, i + 1, config):
          yield config.copy()


def generate_all_combination(param_config):
  keys = sorted(list(param_config.keys()))
  for config in _generate_config(param_config, keys, 0, {}):
    yield config


def get_configs(config_path):
  with open(config_path, 'r') as fi:
    config = json.load(fi)
    return config['param_config'], config['default_config']


def update_config(config):
  machine = config.get('machine', None)
  if machine is not None:
    del config['machine']
    config['feed_machine'] = machine
  if 'expiry0' in config and 'expiry1' in config:
    config['expiry'] = '%s,%s' % (config['expiry0'], config['expiry1'])
    del config['expiry0']
    del config['expiry1']
  return config


def main(argv):
  FLAGS.slurm_job_name_prefix = 'basis_smm_sim'
  assert FLAGS.config_file

  trading_dates = FLAGS.trading_date
  trading_dates_dt = get_trading_dates(FLAGS.trading_date)
  print(trading_dates)
  configs = []
  config_file = FLAGS.config_file
  if ',' in config_file:
    config_file = config_file.split(',')
  else:
    config_file = [config_file]
  for cfile in config_file:
    param_config, default_config = get_configs(cfile)
    assert 'exchanges' in default_config
    assert 'base' in default_config
    assert 'allowed_quote' in default_config
    assert 'sub_dir' in default_config
    if FLAGS.base:
      default_config['base'] = FLAGS.base
    if FLAGS.lot_size_ratio > 0 and 'lot_size_ratio' in default_config:
      default_config['lot_size_ratio'] = FLAGS.lot_size_ratio
    if FLAGS.pos_size_ratio != 0 and 'pos_size_ratio' in default_config:
      default_config['pos_size_ratio'] = FLAGS.pos_size_ratio
    if FLAGS.ref_price_method >= 0:
      default_config['ref_price_method'] = FLAGS.ref_price_method
    if FLAGS.push_pull_method >= 0:
      default_config['push_pull_method'] = FLAGS.push_pull_method
    param_config['trading_date'] = trading_dates_dt
    default_config['run_from_csv'] = FLAGS.run_from_csv
    default_config['output_dir'] = os.path.join(FLAGS.output_dir, default_config['sub_dir'])
    default_config['csv_feed_root'] = FLAGS.csv_feed_root
    default_config['fastfeed'] = FLAGS.fastfeed
    del default_config['sub_dir']
    for config in generate_all_combination(param_config):
      if config['max_std_ratio'] < config['min_std_ratio']:
        continue
      config.update(default_config)
      config = update_config(config)
      configs.append(config)
  run_jobs('basis_smm', 'coin.experimental.yuxuan.basis_smm_strategy', configs)


if __name__ == '__main__':
  flags.DEFINE_float('lot_size_ratio', 0, '')
  flags.DEFINE_float('pos_size_ratio', 1.0, '')
  flags.DEFINE_integer('ref_price_method', -1, '')
  flags.DEFINE_integer('push_pull_method', -1, '')
  app.run(main)
